import xml.etree.ElementTree as ET


def extract(tree, part_name):    # part_name为'tasks' 'models' 'failure_models' or 'repair_models'
    root = tree

    # 找到所有 <models> 元素
    models_elements = root.findall(part_name)

    if not models_elements:
        print("No <models> elements found in the input file.")
        return None

    # 创建新的XML根元素
    new_root = ET.Element('root')

    # 将所有 <models> 元素添加到新的根元素中
    for models in models_elements:
        new_root.append(models)

    return ET.ElementTree(new_root)
